import numpy as np

# *************** Matrix Inversion method for Error Mitigation ***************
def matrix_inversion_mitigation(measured_vector: np.ndarray, response_matrix: np.ndarray) -> np.ndarray:
    """
    Perform matrix inversion for error mitigation.
    
    Args:
        measured_vector: Vector of measured probabilities
        response_matrix: Response matrix R
        
    Returns:
        Mitigated probability vector
    """
    return np.linalg.inv(response_matrix) @ measured_vector

def IBU(ymes: np.ndarray, t0: np.ndarray, Rin: np.ndarray, n: int) -> np.ndarray:
    """
    Perform Iterative Bayesian Unfolding method for error mitigation.
    
    Args:
        ymes: Measured probability distribution
        t0: Initial guess for true distribution
        Rin: Response matrix (measured vs true)
        n: Number of iterations
    
    Returns:
        Mitigated probability distribution
    """
    tn = t0
    for _ in range(n):
        out = np.zeros(t0.shape)
        for j in range(len(t0)):
            mynum = 0.
            for i in range(len(ymes)):
                myden = sum(Rin[i][k] * tn[k] for k in range(len(t0)))
                if myden > 0:
                    mynum += Rin[i][j] * tn[j] * ymes[i] / myden
            out[j] = mynum
        tn = out
    return tn